// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/optional.h>
#include <turbo/log/logging.h>
#include <chrono>
#include <exception>
#include <functional>
#include <stdexcept>
#include <thread>
#include <vector>

#include <pollux/dwio/common/rand_gen.h>
#include <pollux/common/exception/exception.h>

namespace kumo::pollux::dwio::common {
    class retriable_error : public std::runtime_error {
    public:
        explicit retriable_error(const std::exception &e)
            : std::runtime_error(e.what()) {
        }

        explicit retriable_error(const std::string &errorMsg)
            : std::runtime_error(errorMsg) {
        }
    };

    class retries_canceled : public std::runtime_error {
    public:
        explicit retries_canceled()
            : std::runtime_error("Retries have been canceled.") {
        }
    };

    class retries_exhausted : public std::runtime_error {
    public:
        explicit retries_exhausted(const std::exception &e)
            : std::runtime_error(e.what()) {
        }
    };

    using RetryDuration =
    std::chrono::duration<float, std::chrono::milliseconds::period>;

    namespace retrypolicy {
        class IRetryPolicy {
        public:
            virtual ~IRetryPolicy() = default;

            virtual melon::Optional<RetryDuration> nextWaitTime() = 0;

            virtual void start() {
            }
        };

        class KAttempts : public IRetryPolicy {
        public:
            explicit KAttempts(std::vector<RetryDuration> durations)
                : index_(0), durations_(std::move(durations)) {
            }

            melon::Optional<RetryDuration> nextWaitTime() override {
                if (index_ < durations_.size()) {
                    return melon::Optional<RetryDuration>(durations_[index_++]);
                } else {
                    return melon::Optional<RetryDuration>();
                }
            }

        private:
            size_t index_;
            const std::vector<RetryDuration> durations_;
        };

        class ExponentialBackoff : public IRetryPolicy {
        public:
            ExponentialBackoff(
                RetryDuration start,
                RetryDuration max,
                uint64_t maxRetries,
                RetryDuration maxTotal,
                bool countExecutionTime)
                : maxWait_(max),
                  maxTotal_(maxTotal),
                  countExecutionTime_(countExecutionTime),
                  nextWait_(start),
                  totalWait_(0),
                  retriesLeft_(maxRetries) {
                DWIO_ENSURE_LE(start.count(), max.count());
                DWIO_ENSURE(maxTotal_.count() == 0 || maxTotal_.count() > start.count());
            }

            void start() override {
                startTime_ = std::chrono::system_clock::now();
            }

            melon::Optional<RetryDuration> nextWaitTime() override {
                if (retriesLeft_ == 0 || (maxTotal_.count() > 0 && total() >= maxTotal_)) {
                    return melon::Optional<RetryDuration>();
                }

                RetryDuration waitTime = nextWait_ + jitter();
                nextWait_ = std::min(nextWait_ + nextWait_, maxWait_);
                --retriesLeft_;
                totalWait_ += waitTime;
                return melon::Optional<RetryDuration>(waitTime);
            }

        private:
            RetryDuration jitter() {
                return RetryDuration(rand_.gen(melon::to<int64_t>(nextWait_.count()) / 2));
            }

            RetryDuration total() const {
                return countExecutionTime_
                           ? std::chrono::system_clock::now() - startTime_
                           : totalWait_;
            }

            const RetryDuration maxWait_;
            const RetryDuration maxTotal_;
            const bool countExecutionTime_;
            std::chrono::system_clock::time_point startTime_;
            RetryDuration nextWait_;
            RetryDuration totalWait_;
            uint64_t retriesLeft_;
            RandGen rand_;
        };

        class IRetryPolicyFactory {
        public:
            virtual std::unique_ptr<IRetryPolicy> getRetryPolicy() const = 0;

            virtual ~IRetryPolicyFactory() = default;
        };

        class KAttemptsPolicyFactory : public IRetryPolicyFactory {
        private:
            std::vector<RetryDuration> durations_;

        public:
            KAttemptsPolicyFactory(std::vector<RetryDuration> durations)
                : durations_{durations} {
            }

            std::unique_ptr<IRetryPolicy> getRetryPolicy() const {
                return std::make_unique<KAttempts>(durations_);
            }
        };

        class ExponentialBackoffPolicyFactory : public IRetryPolicyFactory {
        public:
            ExponentialBackoffPolicyFactory(
                RetryDuration start,
                RetryDuration maxWait,
                uint64_t maxRetries = std::numeric_limits<uint64_t>::max(),
                RetryDuration maxTotal = RetryDuration::zero(),
                bool countExecutionTime = false)
                : start_(start),
                  maxWait_(maxWait),
                  maxRetries_(maxRetries),
                  maxTotal_(maxTotal),
                  countExecutionTime_(countExecutionTime) {
            }

            std::unique_ptr<IRetryPolicy> getRetryPolicy() const {
                return std::make_unique<ExponentialBackoff>(
                    start_, maxWait_, maxRetries_, maxTotal_, countExecutionTime_);
            }

        private:
            const RetryDuration start_;
            const RetryDuration maxWait_;
            const uint64_t maxRetries_;
            const RetryDuration maxTotal_;
            const bool countExecutionTime_;
        };
    } // namespace retrypolicy

    class RetryModule {
    public:
        template<typename F>
        static auto withRetry(
            F func,
            std::unique_ptr<retrypolicy::IRetryPolicy> policy,
            std::function<bool()> abortFunc = nullptr) {
            policy->start();
            do {
                try {
                    // If abort signal is triggered before func, no ops.
                    if (abortFunc && abortFunc()) {
                        KLOG(WARNING) << "Retries canceled.";
                        throw retries_canceled();
                    }

                    return func();
                } catch (const retriable_error &error) {
                    KLOG(INFO) << "RetryModule caught retriable exception. " << error.what();
                    melon::Optional<RetryDuration> wait = policy->nextWaitTime();
                    if (wait.has_value()) {
                        auto ms = wait.value().count();
                        KLOG(INFO) << "RetryModule : Waiting for " << ms
                    << "(ms) before retry";
                        while (ms > 0) {
                            // check abort signal every second.
                            if (abortFunc && abortFunc()) {
                                KLOG(WARNING) << "Retries canceled.";
                                throw retries_canceled();
                            }

                            const float abortCheckInterval = 1000;
                            std::this_thread::sleep_for(
                                RetryDuration(std::min(ms, abortCheckInterval)));
                            ms -= abortCheckInterval;
                        }
                    } else {
                        KLOG(WARNING) << "Retries exhausted!!!";
                        throw retries_exhausted(error);
                    }
                }
            } while (true);
        }
    };
} // namespace kumo::pollux::dwio::common
